import datetime
import pandas
import peewee
import sqlite3
from typing import Any, Dict, List, Set, Tuple, Type, Optional, Union, Callable
import vnpy.tools.chinese_mainland_time_config
import vnpy.tools.utility
import vnpy.trader.database
from vnpy.trader.constant import Interval, Exchange


class DataMissingChecker(object):
    """数据缺失检查器"""

    @classmethod
    def get_peewee_SqliteDatabase(cls) -> Optional[peewee.SqliteDatabase]:
        """"""
        sqlite_database: Optional[peewee.SqliteDatabase] = None

        import vnpy_sqlite

        database: vnpy.trader.database.BaseDatabase = vnpy.trader.database.get_database()

        if isinstance(database, vnpy_sqlite.sqlite_database.SqliteDatabase):
            database: vnpy_sqlite.sqlite_database.SqliteDatabase = database

            if isinstance(database.db, peewee.SqliteDatabase):
                sqlite_database: peewee.SqliteDatabase = database.db

        return sqlite_database

    @classmethod
    def fetchall(cls, cursor: sqlite3.Cursor) -> List[dict]:
        """"""
        head = [col[0] for col in cursor.description]
        rows = [dict(zip(head, data)) for data in cursor.fetchall()]
        return rows

    @classmethod
    def select_count(
        cls,
        interval: Interval,
        exchange: Exchange,
        symbol: str,
        business_day_ge: Optional[datetime.date] = None,
        business_day_le: Optional[datetime.date] = None,
    ) -> List[dict]:
        """"""
        and_business_day_ge: str = ""
        and_business_day_le: str = ""

        if business_day_ge:
            and_business_day_ge = f""" AND business_day>='{business_day_ge.strftime("%Y-%m-%d")}' """
        if business_day_le:
            and_business_day_le = f""" AND business_day<='{business_day_le.strftime("%Y-%m-%d")}' """

        sql_statement: str = f"""
SELECT
    interval,
    exchange,
    symbol,
    CASE
        WHEN (strftime('%H',datetime)+0 >=  7 AND strftime('%H',datetime)+0 < 19) THEN DATE(datetime) -- 日盘
        WHEN  strftime('%H',datetime)+0 >= 19 THEN DATE(datetime, '+' || (CASE strftime('%w',datetime)+0 WHEN 5 THEN 3 WHEN 6 THEN 2 ELSE 1 END) || ' day') -- 夜盘_半夜
        WHEN  strftime('%H',datetime)+0 <   7 THEN DATE(datetime, '+' || (CASE strftime('%w',datetime)+0 WHEN 6 THEN 2 WHEN 7 THEN 1 ELSE 0 END) || ' day') -- 夜盘_凌晨
        ELSE DATE(datetime) END AS business_day,
    COUNT(*) AS count
FROM dbbardata WHERE interval='{interval.value}' AND exchange='{exchange.value}' AND symbol='{symbol}' {and_business_day_ge}{and_business_day_le}
    GROUP BY interval, exchange, symbol, business_day;
    """

        sqlite_database: peewee.SqliteDatabase = cls.get_peewee_SqliteDatabase()

        cursor = sqlite_database.execute_sql(sql=sql_statement)
        assert isinstance(cursor, sqlite3.Cursor)

        rows: List[dict] = cls.fetchall(cursor=cursor)

        return rows

    @classmethod
    def diff_count(
        cls,
        interval: Interval,
        exchange: Exchange,
        symbol: str,
        business_day_ge: Union[datetime.date, str, None] = '2020-05-01',
        business_day_le: Union[datetime.date, str, None] = None,
    ) -> None:
        """"""
        if isinstance(business_day_ge, str):
            business_day_ge: datetime.date = datetime.datetime.strptime(business_day_ge, "%Y-%m-%d").date()
        if isinstance(business_day_le, str):
            business_day_le: datetime.date = datetime.datetime.strptime(business_day_le, "%Y-%m-%d").date()

        rows: List[dict] = cls.select_count(
            interval=interval,
            exchange=exchange,
            symbol=symbol,
            business_day_ge=business_day_ge,
            business_day_le=business_day_le,
        )

        date_2_count_db: Dict[datetime.date, int] = {datetime.datetime.strptime(_["business_day"], "%Y-%m-%d").date(): _["count"] for _ in rows}

        min_date: datetime.date = min(date_2_count_db.keys())
        max_date: datetime.date = max(date_2_count_db.keys())

        date_2_count_py: Dict[datetime.date, int] = {}

        for idx in range(0, (max_date - min_date).days + 1, 1):
            cur_date: datetime.date = min_date + datetime.timedelta(days=idx)

            dttm_s: Optional[List[datetime.datetime]] = vnpy.tools.chinese_mainland_time_config.ChineseMainlandTimeConfig.calc_for_business_day(
                exchange=exchange,
                symbol=symbol,
                unit_second=60,
                business_day=cur_date,
            )
            if not dttm_s:
                continue

            date_2_count_py[cur_date] = len(dttm_s) - 1  # 数据里附加了"这个业务日的 结束时间戳", 这里需要去掉它,

        lines = []
        date_s: List[datetime.date] = sorted(set(date_2_count_db.keys()).union(date_2_count_py.keys()))
        for date in date_s:
            line = {
                "date": date,
                "count_db": date_2_count_db.get(date),
                "count_py": date_2_count_py.get(date),
                "interval": interval.value,
                "exchange": exchange.value,
                "symbol": symbol,
            }
            lines.append(line)

        df1 = pandas.DataFrame(data=lines)
        df1.set_index("date", inplace=True, verify_integrity=False)
        df2 = df1[df1["count_db"] != df1["count_py"]]

        pandas.set_option('display.max_rows', None)  # 显示所有行
        print(df2)

    @classmethod
    def select_datetime(
        cls,
        interval: Interval,
        exchange: Exchange,
        symbol: str,
        business_day: datetime.date,
    ) -> List[dict]:
        """"""
        sql_statement: str = f"""
    SELECT
        interval,
        exchange,
        symbol,
        CASE
            WHEN (strftime('%H',datetime)+0 >=  7 AND strftime('%H',datetime)+0 < 19) THEN DATE(datetime) -- 日盘
            WHEN  strftime('%H',datetime)+0 >= 19 THEN DATE(datetime, '+' || (CASE strftime('%w',datetime)+0 WHEN 5 THEN 3 WHEN 6 THEN 2 ELSE 1 END) || ' day') -- 夜盘_半夜
            WHEN  strftime('%H',datetime)+0 <   7 THEN DATE(datetime, '+' || (CASE strftime('%w',datetime)+0 WHEN 6 THEN 2 WHEN 7 THEN 1 ELSE 0 END) || ' day') -- 夜盘_凌晨
            ELSE DATE(datetime) END AS business_day,
        datetime
    FROM dbbardata WHERE interval='{interval.value}' AND exchange='{exchange.value}' AND symbol='{symbol}' AND business_day='{business_day.strftime("%Y-%m-%d")}';
    """

        sqlite_database: peewee.SqliteDatabase = cls.get_peewee_SqliteDatabase()

        cursor = sqlite_database.execute_sql(sql=sql_statement)
        assert isinstance(cursor, sqlite3.Cursor)

        rows: List[dict] = cls.fetchall(cursor=cursor)

        return rows

    @classmethod
    def diff_dttm(
        cls,
        interval: Interval,
        exchange: Exchange,
        symbol: str,
        business_day: Union[datetime.date, str],
    ) -> None:
        """"""
        if isinstance(business_day, str):
            business_day: datetime.date = datetime.datetime.strptime(business_day, "%Y-%m-%d").date()

        rows: List[dict] = cls.select_datetime(
            interval=interval,
            exchange=exchange,
            symbol=symbol,
            business_day=business_day,
        )

        dttm_2_exist_db: Dict[datetime.datetime, bool] = {datetime.datetime.strptime(_["datetime"], "%Y-%m-%d %H:%M:%S"): True for _ in rows}

        dttm_s: Optional[List[datetime.datetime]] = vnpy.tools.chinese_mainland_time_config.ChineseMainlandTimeConfig.calc_for_business_day(
            exchange=exchange,
            symbol=symbol,
            unit_second=60,
            business_day=business_day,
        )

        dttm_2_exist_py: Dict[datetime.datetime, bool] = {_: True for _ in dttm_s[:-1]}

        lines = []
        dttm_s: List[datetime.date] = sorted(set(dttm_2_exist_db.keys()).union(dttm_2_exist_py.keys()))
        for dttm in dttm_s:
            line = {
                "dttm": dttm,
                "exist_db": dttm_2_exist_db.get(dttm),
                "exist_py": dttm_2_exist_py.get(dttm),
                "interval": interval.value,
                "exchange": exchange.value,
                "symbol": symbol,
                "business_day": business_day,
            }
            lines.append(line)

        df1 = pandas.DataFrame(data=lines)
        df1.set_index("dttm", inplace=True, verify_integrity=False)
        df2 = df1[df1["exist_db"] != df1["exist_py"]]

        pandas.set_option('display.max_rows', None)  # 显示所有行
        print(df2)


if __name__ == '__main__':
    if 1:
        DataMissingChecker.diff_count(
            interval=Interval.MINUTE,
            exchange=Exchange.SHFE,
            symbol="AGL8",
        )
    if 1:
        DataMissingChecker.diff_dttm(
            interval=Interval.MINUTE,
            exchange=Exchange.SHFE,
            symbol="AGL8",
            business_day="2023-09-06",
        )
